'''
Plots unstructured samples with 2 locations.

@author: Per Rosengren
'''
import numpy as np
import OpenGL.GL as gl
import glumpy
import scipy.spatial

class Unstructured2d(object):
    def __init__(self, x, y, z, 
            data_lim=None, 
            dir_x_right=True, 
            dir_y_top=True, 
            colormap=glumpy.colormap.IceAndFire, frame=None):
        self._create_shader(colormap)
        self._create_mesh(x, y, z)
        if data_lim is not None:
            self._data_lim = data_lim
        else:
            self._data_lim = [[x.min(), x.max()], [y.min(), y.max()]]
        self._dir_x_right = dir_x_right
        self._dir_y_top = dir_y_top
        self._frame = frame
        self._frame.push(self)
        
    def _create_mesh(self, x, y, z):
        tri = scipy.spatial.Delaunay(np.column_stack([x, y]))
        I = tri.simplices
        t = z.astype('f4')
        t = .5 + .5 * t / np.abs(t).max()
        V = np.column_stack((
                x.astype('f4'), 
                y.astype('f4'), 
                t)
            ).view(
            dtype = [('position', 'f4', 2),
                     ('tex_coord', 'f4', t.shape[1] if len(t.shape) > 1 else 1),
                    ])
        self._mesh = glumpy.graphics.VertexBuffer(V, I)

    def _create_shader(self, colormap):
        self._color_lut = glumpy.image.Texture(colormap)
        vertex = """
            void main()
            {
                gl_FrontColor = gl_Color;
                gl_TexCoord[0] = gl_MultiTexCoord0;
                gl_Position = ftransform();
            }"""
    
        fragment_t = """
            uniform sampler{0}D color_lut; /* default location 2 */
            
            void main() 
            {{ 
                gl_FragColor = texture{0}D(color_lut,gl_TexCoord[0].{1});;
            }}"""
        if self._color_lut.height == 0:
            fragment = fragment_t.format("1", "x")
        else:
            fragment = fragment_t.format("2", "xy")
        self._shader = glumpy.graphics.Shader(vertex,fragment)
    
    def on_resize(self, width, height, x=0, y=0):
        data_width = self._data_lim[0][1] - self._data_lim[0][0]
        data_height = self._data_lim[1][1] - self._data_lim[1][0]
        data_aspect = data_width / float(data_height)
        frame_aspect = width / float(height)
        if frame_aspect >= data_aspect:
            padding = (frame_aspect * data_height - data_width) / 2.
            frame_lim = [
                [self._data_lim[0][0] - padding, 
                 self._data_lim[0][1] + padding],
                [self._data_lim[1][0], 
                 self._data_lim[1][1]]] 
        else:
            padding = (data_width / frame_aspect - data_height) / 2.
            frame_lim = [
                [self._data_lim[0][0], 
                 self._data_lim[0][1]],
                [self._data_lim[1][0] - padding, 
                 self._data_lim[1][1] + padding]] 
        self._ortho = frame_lim[0][::(1 if self._dir_x_right else -1)]
        self._ortho += frame_lim[1][::(1 if self._dir_y_top else -1)]
        self._ortho += -1000, 1000   

    def _projection_push(self):
        gl.glMatrixMode(gl.GL_PROJECTION)
        gl.glPushMatrix()
        gl.glLoadIdentity( )
        gl.glOrtho(*self._ortho)

    def _projection_pop(self):
        gl.glMatrixMode(gl.GL_PROJECTION)
        gl.glPopMatrix()

    def _shader_push(self):
        self._shader.bind()
        gl.glActiveTexture(gl.GL_TEXTURE2)
        gl.glBindTexture(self._color_lut.target, self._color_lut.texture)
        self._shader.uniformi('color_lut', 2)

    def _shader_pop(self):
        self._shader.unbind()
        
    def on_draw(self):
        self._frame.lock()
        self._frame.draw()
        self._shader_push()
        self._projection_push()
        self._mesh.draw(gl.GL_TRIANGLES)
        self._projection_pop()
        self._shader_pop()
        self._frame.unlock()

def create_colormap2d_hsv(size=512):
    import matplotlib.colors
    import math
    u, v = np.meshgrid(np.linspace(-1, 1, size), np.linspace(-1, 1, size))
    hsv = np.ones((size, size, 3), dtype=np.float32)
    hsv[:, :, 0] = (np.arctan2(u, v) / (2*math.pi) + .5)
    hsv[:, :, 1] = np.minimum(1., np.sqrt(u**2 + v**2))
    rgb = matplotlib.colors.hsv_to_rgb(hsv)
    return rgb

def create_colormap2d_4dirs(size=512):
    rgb = np.ones((size, size, 3), dtype=np.float32)
    hs = size / 2
    u, v = np.meshgrid(np.linspace(1, 0, hs), np.linspace(1, 0, hs))
    rgb[:hs, :hs, 0] = 1.
    rgb[:hs, :hs, 1] = 1. - v + u/2.
    rgb[:hs, :hs, 2] = 1. - np.maximum(u, v)
    u = u[:, ::-1]
    rgb[:hs, hs:, 0] = 1. - u + v
    rgb[:hs, hs:, 1] = 1. - np.maximum(u, v)
    rgb[:hs, hs:, 2] = 1. - v + u
    v = v[::-1, :]
    rgb[hs:, hs:, 0] = 1. - np.maximum(u, v)
    rgb[hs:, hs:, 1] = 1. - u + v
    rgb[hs:, hs:, 2] = 1. - v + u
    u = u[:, ::-1]
    rgb[hs:, :hs, 0] = 1. - v + u/2.
    rgb[hs:, :hs, 1] = 1.
    rgb[hs:, :hs, 2] = 1. - np.maximum(u, v)
    rgb = np.minimum(1., rgb)
    return rgb

if __name__ == '__main__':

    fig_main = glumpy.figure(size=(800,400), position=(100, 100))
    subfigs = fig_main.split(position='horizontal')
    @fig_main.event
    def on_draw():
        fig_main.clear(0.90,0.90,0.90,1.00)

    
#     @frame.event
#     def on_draw():
#         frame.lock()
#         frame.draw()
#         samples.draw()
#         frame.unlock()

    loc = np.meshgrid(np.linspace(0, 1, 100), np.linspace(0, 1, 100))
    loc = np.column_stack([v.flat for v in loc])
    np.random.shuffle(loc)
    v = loc - .5
    samples = Unstructured2d(loc[:, 0], loc[:, 1], v, colormap=create_colormap2d_4dirs(size=128), frame=subfigs[0].add_frame())
    v = loc[:, 0] - .5
    samples = Unstructured2d(loc[:, 0], loc[:, 1], v, colormap=glumpy.colormap.IceAndFire.LUT['rgb'][1:].flatten().view((np.float32,3)), frame=subfigs[1].add_frame())
    fig_main.show()
